# Copyright 2025 the LlamaFactory team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from abc import abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Optional, Union

from ..extras import logging
from .data_utils import Role


if TYPE_CHECKING:
    from datasets import Dataset, IterableDataset
    from transformers import Seq2SeqTrainingArguments

    from ..hparams import DataArguments
    from .mm_plugin import AudioInput, ImageInput, VideoInput
    from .parser import DatasetAttr

    MediaType = Union[ImageInput, VideoInput, AudioInput]


logger = logging.get_logger(__name__)


@dataclass
class DatasetConverter:
    dataset_attr: "DatasetAttr"
    data_args: "DataArguments"

    # 查找媒体文件路径，如果从本地加载则拼接 media_dir
    def _find_medias(self, medias: Union["MediaType", list["MediaType"], None]) -> Optional[list["MediaType"]]:
        r"""Optionally concatenate media path to media dir when loading from local disk."""
        # 如果没有提供媒体信息，返回 None
        if medias is None:
            return None
        # 如果不是列表类型，则包装成列表
        elif not isinstance(medias, list):
            medias = [medias]
        # 如果是空列表，也返回 None
        elif len(medias) == 0:
            return None
        else:
            # 复制一份避免修改原始数据
            medias = medias[:]

        # 如果是从本地脚本或文件加载，并且第一个元素是字符串路径
        if self.dataset_attr.load_from in ["script", "file"] and isinstance(medias[0], str):
            for i in range(len(medias)):
                # 检查 media_dir 下是否存在该路径，若存在则拼接 media_dir
                if os.path.isfile(os.path.join(self.data_args.media_dir, medias[i])):
                    medias[i] = os.path.join(self.data_args.media_dir, medias[i])
                else:
                    # 否则保留原路径并打印警告
                    logger.warning_rank0_once(f"Media {medias[i]} does not exist in `media_dir`. Use original path.")

        # 返回处理后的媒体路径列表
        return medias

    # 抽象方法：用于将单个样本转换为标准格式，需要子类重写
    @abstractmethod
    def __call__(self, example: dict[str, Any]) -> dict[str, Any]:
        r"""Convert a single example in the dataset to the standard format."""
        ...


@dataclass
class AlpacaDatasetConverter(DatasetConverter):
    def __call__(self, example: dict[str, Any]) -> dict[str, Any]:
        # 初始化空列表用于构建 prompt 对话历史
        prompt = []
        # 如果存在历史对话字段且是列表类型，则逐条转换为标准格式
        if self.dataset_attr.history and isinstance(example[self.dataset_attr.history], list):
            for old_prompt, old_response in example[self.dataset_attr.history]:
                # 添加用户提问到 prompt
                prompt.append({"role": Role.USER.value, "content": old_prompt})
                # 添加助手回复到 prompt
                prompt.append({"role": Role.ASSISTANT.value, "content": old_response})

        # 初始化 query 列表，用于拼接 prompt 和 query 字段内容
        query = []
        # 如果存在 prompt 字段且有值，加入 query 列表
        if self.dataset_attr.prompt and example[self.dataset_attr.prompt]:
            query.append(example[self.dataset_attr.prompt])

        # 如果存在 query 字段且有值，也加入 query 列表
        if self.dataset_attr.query and example[self.dataset_attr.query]:
            query.append(example[self.dataset_attr.query])

        # 将合并后的 prompt + query 作为新的用户输入添加进 prompt
        prompt.append({"role": Role.USER.value, "content": "\n".join(query)})  # "prompt\nquery"

        # 如果是 KTO 模式，并且标记为布尔值，则构造偏好响应对
        if self.dataset_attr.kto_tag and isinstance(example[self.dataset_attr.kto_tag], bool):  # kto example
            # 构造基础 response
            response = [{"role": Role.ASSISTANT.value, "content": example[self.dataset_attr.response]}]
            # 正向偏好：在 response 后加空回复
            if example[self.dataset_attr.kto_tag]:
                response = response + [{"role": Role.ASSISTANT.value, "content": ""}]
            # 反向偏好：在 response 前加空回复
            else:
                response = [{"role": Role.ASSISTANT.value, "content": ""}] + response
        # 如果是排序任务（ranking），并且 chosen 和 rejected 是字符串类型
        elif (
            self.dataset_attr.ranking
            and isinstance(example[self.dataset_attr.chosen], str)
            and isinstance(example[self.dataset_attr.rejected], str)
        ):  # pairwise example
            # 构造两个候选 response
            response = [
                {"role": Role.ASSISTANT.value, "content": example[self.dataset_attr.chosen]},
                {"role": Role.ASSISTANT.value, "content": example[self.dataset_attr.rejected]},
            ]
        # 如果是正常响应字段且为字符串类型
        elif self.dataset_attr.response and isinstance(example[self.dataset_attr.response], str):  # normal example
            # 构造单个 response
            response = [{"role": Role.ASSISTANT.value, "content": example[self.dataset_attr.response]}]
        # 无有效 response 的情况（如未标注数据）
        else:  # unsupervised
            response = []

        # 最终输出统一格式的数据结构
        output = {
            "_prompt": prompt,
            "_response": response,
            "_system": example[self.dataset_attr.system] if self.dataset_attr.system else "",
            "_tools": example[self.dataset_attr.tools] if self.dataset_attr.tools else "",
            "_images": self._find_medias(example[self.dataset_attr.images]) if self.dataset_attr.images else None,
            "_videos": self._find_medias(example[self.dataset_attr.videos]) if self.dataset_attr.videos else None,
            "_audios": self._find_medias(example[self.dataset_attr.audios]) if self.dataset_attr.audios else None,
        }
        return output


@dataclass
class SharegptDatasetConverter(DatasetConverter):
    def __call__(self, example: dict[str, Any]) -> dict[str, Any]:
        # 定义原始角色标签到标准 Role 枚举值的映射关系
        tag_mapping = {
            self.dataset_attr.user_tag: Role.USER.value,
            self.dataset_attr.assistant_tag: Role.ASSISTANT.value,
            self.dataset_attr.observation_tag: Role.OBSERVATION.value,
            self.dataset_attr.function_tag: Role.FUNCTION.value,
            self.dataset_attr.system_tag: Role.SYSTEM.value,
        }
        # 定义奇数轮应出现的角色标签（如 user、observation）
        odd_tags = (self.dataset_attr.user_tag, self.dataset_attr.observation_tag)
        # 定义偶数轮应出现的角色标签（如 assistant、function）
        even_tags = (self.dataset_attr.assistant_tag, self.dataset_attr.function_tag)
        # 奇偶轮次允许的标签组合
        accept_tags = (odd_tags, even_tags)
        # 获取消息列表字段
        messages = example[self.dataset_attr.messages]
        # 如果存在 system_tag 并且第一条是 system 角色，则提取 system 内容并去掉该条目
        if (
            self.dataset_attr.system_tag
            and len(messages) != 0
            and messages[0][self.dataset_attr.role_tag] == self.dataset_attr.system_tag
        ):
            system = messages[0][self.dataset_attr.content_tag]
            messages = messages[1:]
        else:
            # 否则从系统字段中获取 system 内容
            system = example[self.dataset_attr.system] if self.dataset_attr.system else ""

        # 初始化对齐后的消息列表
        aligned_messages = []
        # 标记数据是否异常
        broken_data = False
        # 遍历每一轮对话，检查格式是否正确并进行标准化
        for turn_idx, message in enumerate(messages):
            # 检查当前轮次的角色是否符合预期（奇偶交替）
            if message[self.dataset_attr.role_tag] not in accept_tags[turn_idx % 2]:
                logger.warning_rank0(f"Invalid role tag in {messages}.")
                broken_data = True
                break

            # 将角色和内容标准化后加入列表
            aligned_messages.append(
                {
                    "role": tag_mapping[message[self.dataset_attr.role_tag]],
                    "content": message[self.dataset_attr.content_tag],
                }
            )

        # 检查对话是否为成对的问答结构（非 ranking 情况下应为偶数条消息）
        if (not self.dataset_attr.ranking and len(aligned_messages) % 2 != 0) or (
            self.dataset_attr.ranking and len(aligned_messages) % 2 == 0
        ):
            logger.warning_rank0(f"Invalid message count in {messages}.")
            broken_data = True

        # 如果数据异常，跳过该样本或构造空 prompt/response
        if broken_data:
            logger.warning_rank0("Skipping this abnormal example.")
            prompt, response = [], []
        # 如果是 KTO 模式，根据偏好构造 prompt 和 response
        elif self.dataset_attr.kto_tag and isinstance(example[self.dataset_attr.kto_tag], bool):  # kto example
            prompt = aligned_messages[:-1]
            response = aligned_messages[-1:]
            # 正向偏好：在 response 后加空回复
            if example[self.dataset_attr.kto_tag]:
                response = response + [{"role": Role.ASSISTANT.value, "content": ""}]
            # 反向偏好：在 response 前加空回复
            else:
                response = [{"role": Role.ASSISTANT.value, "content": ""}] + response
        # 如果是排序任务（ranking），处理 chosen 和 rejected 回复
        elif (
            self.dataset_attr.ranking
            and isinstance(example[self.dataset_attr.chosen], dict)
            and isinstance(example[self.dataset_attr.rejected], dict)
        ):  # pairwise example
            chosen = example[self.dataset_attr.chosen]
            rejected = example[self.dataset_attr.rejected]
            # 检查 chosen 和 rejected 的角色是否合法
            if (
                chosen[self.dataset_attr.role_tag] not in accept_tags[-1]
                or rejected[self.dataset_attr.role_tag] not in accept_tags[-1]
            ):
                logger.warning_rank0(f"Invalid role tag in {[chosen, rejected]}.")
                broken_data = True

            # 构造 prompt 为完整的对话历史
            prompt = aligned_messages
            # 构造 response 为两个候选回复
            response = [
                {
                    "role": tag_mapping[chosen[self.dataset_attr.role_tag]],
                    "content": chosen[self.dataset_attr.content_tag],
                },
                {
                    "role": tag_mapping[rejected[self.dataset_attr.role_tag]],
                    "content": rejected[self.dataset_attr.content_tag],
                },
            ]
        else:  # normal example
            # 正常情况：prompt 是除最后一条外的所有消息，response 是最后一条
            prompt = aligned_messages[:-1]
            response = aligned_messages[-1:]

        # 最终输出统一格式的数据结构
        output = {
            "_prompt": prompt,
            "_response": response,
            "_system": system,
            "_tools": example[self.dataset_attr.tools] if self.dataset_attr.tools else "",
            "_images": self._find_medias(example[self.dataset_attr.images]) if self.dataset_attr.images else None,
            "_videos": self._find_medias(example[self.dataset_attr.videos]) if self.dataset_attr.videos else None,
            "_audios": self._find_medias(example[self.dataset_attr.audios]) if self.dataset_attr.audios else None,
        }
        return output


# 数据组装的格式,就两种格式.
DATASET_CONVERTERS = {
    "alpaca": AlpacaDatasetConverter,
    "sharegpt": SharegptDatasetConverter,
}


# 注册一个新的数据集转换器，用于将不同格式的数据统一成标准结构
def register_dataset_converter(name: str, dataset_converter: type["DatasetConverter"]) -> None:
    r"""Register a new dataset converter."""
    # 如果该名称的转换器已存在，则抛出异常防止重复注册
    if name in DATASET_CONVERTERS:
        raise ValueError(f"Dataset converter {name} already exists.")

    # 将指定名称与对应的转换器类注册到全局字典中
    DATASET_CONVERTERS[name] = dataset_converter


# 获取一个已注册的数据集转换器实例
def get_dataset_converter(name: str, dataset_attr: "DatasetAttr", data_args: "DataArguments") -> "DatasetConverter":
    r"""Get a dataset converter."""
    # 如果请求的转换器未注册，则抛出异常
    if name not in DATASET_CONVERTERS:
        raise ValueError(f"Dataset converter {name} not found.")

    # 使用提供的参数创建并返回对应的数据集转换器实例
    return DATASET_CONVERTERS[name](dataset_attr, data_args)


# 对齐数据集
def align_dataset(
    dataset: Union["Dataset", "IterableDataset"],
    dataset_attr: "DatasetAttr",
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
) -> Union["Dataset", "IterableDataset"]:
    r"""Align the dataset to a specific format.

    Aligned dataset:
    _prompt: [{"role": "user", "content": "..."}] * (2T - 1)
    _response: [{"role": "assistant", "content": "..."}] * N (N > 1 for ranking dataset)
    _system: "..."
    _tools: "..."
    _images: []
    _videos: []
    _audios: []
    """
    # 使用迭代器
    # 获取当前数据集的第一条样本的所有字段名，用于后续删除原始列
    column_names = list(next(iter(dataset)).keys())
    # 初始化 map 函数的参数字典
    kwargs = {}
    # 不使用流式加载数据的话,data_args.preprocessing_num_workers用这个参数
    # data_args.preprocessing_num_workers就是多少个线程加载
    if not data_args.streaming:
        kwargs = dict(
            num_proc=data_args.preprocessing_num_workers,
            load_from_cache_file=(not data_args.overwrite_cache) or (training_args.local_process_index != 0),
            desc="Converting format of dataset",
        )

    # 根据指定的格式化方式获取对应的数据集转换器
    dataset_converter = get_dataset_converter(dataset_attr.formatting, dataset_attr, data_args)
    # 使用 map 方法将每个样本通过转换器统一格式，并删除原始列
    return dataset.map(
        dataset_converter,
        batched=False,
        remove_columns=column_names,
        **kwargs,
    )
